from lxml import etree, objectify
from lxml.etree import ElementTree

class Project:
    def __init__(self, imgPath, xmlPath, imgName ,dt):
        self.imgPath = imgPath
        self.imgName = imgName
        self.xmlPath = xmlPath
        self.dt = dt
        self.e = objectify.ElementMaker(annotate=False)
        #self.labelNames = []

    def genXml(self):
        if self.imgPath is None or \
                self.imgName is None or \
                self.xmlPath is None:
            return None
        imgLabeled_tree = self.e.imgsLabeled(
            self.e.imgName(self.imgName),
            self.e.imgPath(self.imgPath),
            self.e.xmlPath(self.xmlPath)
        )
        return imgLabeled_tree

    def prettify(self, elem):
        """
            Return a pretty-printed XML string for the Element.
        """
        rough_string = ElementTree.tostring(elem, 'utf8')
        root = etree.fromstring(rough_string)
        return etree.tostring(root, pretty_print=True, encoding=ENCODE_METHOD).replace("  ".encode(), "\t".encode())

    def save(self, targetFile=None):
        imgLabeled_tree = self.genXml()
        # if targetFile is None:
        #     out_file = codecs.open(
        #         self.xmlPath + XML_EXT, 'w', encoding=ENCODE_METHOD)
        # else:
        #     out_file = codecs.open(targetFile, 'w', encoding=ENCODE_METHOD)
        #prettifyResult = self.prettify(root)
        #out_file.write(prettifyResult.decode('utf8'))
        #out_file.close()

        filepath = 'imgLabeled' + '_' + self.dt + '.xml'
        f = open(filepath,'a')
        etree.ElementTree(imgLabeled_tree).write(f, pretty_print=True)
        f.close()